import numpy as np

from ..model import Vehicle
from ..sensor import RangeBearingSensor
from ..map import Map


class SLAMExtendedKalmanFilter:
    def __init__(self, model: Vehicle, senor: RangeBearingSensor, map_: Map, P_est: np.ndarray, V_est: np.ndarray,
                 W_est: np.ndarray) -> None:
        super().__init__()

        self._model = model
        self._sensor = senor
        self._map = map_

        self._P_est = P_est
        self._V_est = V_est
        self._W_est = W_est

    def step(self, v, z_measure):

        xv_est = self._model.x
        x_pred = self._model.handle_state(self._model.F(xv_est, v))
        Fx = self._model.Fx(xv_est, v)
        Fv = self._model.Fv(xv_est, v)
        P_pred = Fx @ self._P_est[:3, :3] @ Fx.T + Fv @ self._V_est @ Fv.T
        self._P_est[:3, :3] = P_pred

        for landmark_id, iz in z_measure:
            z_meas = iz[:2]

            if not self._map.is_in(landmark_id):
                g = self._sensor.G(x_pred, z_meas)

                self.add_P_est(x_pred, g, z_meas)

                self._map.add(landmark_id, g)

            lm_idx = self._map.get_index(landmark_id)
            g = self._map.get_feature(landmark_id)

            landmark_est = self._map.get_features().reshape(1, -1).flatten()

            z_pred = self._sensor.H(x_pred, g)
            hh = np.zeros((2, len(landmark_est) + 3))
            hh[:, :3] = self._sensor.Hxv(x_pred, g)
            hh[:, 2 * lm_idx + 3: 2 * (lm_idx + 1) + 3] = self._sensor.Hxi(x_pred, g)

            S = hh @ self._P_est @ hh.T + self._W_est
            K = self._P_est @ hh.T @ np.linalg.inv(S)

            y = self._sensor.handle_measure(z_meas - z_pred)
            delta = K @ y
            x_pred[:] = self._model.handle_state(x_pred + delta[:3])
            landmark_est[:] = landmark_est + delta[3:]

            self._map.update(landmark_est.reshape(-1, 2))
            self._P_est[:, :] = self._P_est - K @ hh @ self._P_est


        self._model.x = x_pred


    def add_P_est(self, xv, xi, z):
        n = self._P_est.shape[0]
        Yz = np.zeros((n + 2, n + 2))
        Yz[:n, :n] = np.eye(n)
        Yz[n:, :3] = self._sensor.Gxv(xv, z)
        Yz[-2:, -2:] = self._sensor.Gz(xv, z)
        P_1 = np.zeros((n + 2, n + 2))
        P_1[:n, :n] = self._P_est
        P_1[-2:, -2:] = self._W_est
        # self._P_est = Yz @ P_1 @ Yz.T
        self._P_est = P_1
